from qutip import *
from nv_system import *
from numpy import *
from rotating_frame import *


class pulsecontrol:
    def __init__(self,pulses,expected_operators=None):
        self.pulse_channels={}
        if not isinstance(pulses,list):
            self.pulses = [pulses]
        else:
            self.pulses = pulses
        self.start_time_index = []
        start_time = 0
        if len(self.pulses)>1:
            for i in range(len(self.pulses)):
                self.start_time_index.append(start_time)
                start_time = start_time + self.pulses[i].get_duration()
        self.duration = start_time
        self.expected_operators = expected_operators
        self.pointer = 0
    
    
    def get_shape(self,channame,t):
        
        if self.duration != 0 :
            if t > self.duration:
                return 0.0,0.0
        else:
            return self.pulses[0].get_shape(channame,t)
        
        while t < self.start_time_index[self.pointer]:
            self.pointer = self.pointer + 1
        while t >= self.start_time_index[self.pointer] +self.pulses[self.pointer].get_duration():
            self.pointer = self.pointer - 1
        
        return self.pulses[self.pointer].get_shape(channame,t) 
            
        
            
def gen_lab_frame_setting(static_ham,control_matrix,mwchannels):    
    Ht = []
    Ht.append(static_ham * 2*pi)    
    matx,maty,matz = control_matrix
    Ht.append([-matx*pi*2,Bdx])
    Ht.append([-matx*pi*0,Bdx])
#     print'waring:by and bz is not included'
    Ht.append([-maty*pi*2,Bdy])
    Ht.append([-matz*pi*2,Bdz])           
    # define the time-dependence of the hamiltonian using the list-string format   
    opts=Odeoptions()
    opts.nsteps=1e6
    args={'MWCH':mwchannels}    
    return args,opts,Ht
    

def gen_rotating_frame_setting(sysdef,rframe,Hleft,channels,s_coup,s_mw,threshold = 1e-6,debug=True):    
    Ht = []
    rh0 = 0
    max_freq = 0
    diagonal_ham_items = []
    nondiagonal_ham_items = []
    import StringIO
    
    sperturb = Perturbation(Hleft,0,1.0)
    diagonal_item,nondiagonal_items = \
        sperturb.rotating_wave_approximate(rframe.P,rframe.omega,s_coup)               
    rh0 = rh0 + diagonal_item
    
    for nondiagonal_item in nondiagonal_items:
        if(abs(nondiagonal_item[1])<threshold):
            rh0 = rh0+nondiagonal_item[0]
        else:
            nondiagonal_item.append('static_item')
            nondiagonal_ham_items.append(nondiagonal_item)
        
    
    for chname in channels:
        channel = channels[chname]
        mwperturb=MWFieldPerturbation(channel)            
        diagonal_item,nondiagonal_items = \
            mwperturb.rotating_wave_approximate(rframe.P,rframe.omega,s_mw)

        if diagonal_item.norm()!=0:
            diagonal_ham_items.append([diagonal_item,chname])
        
        for nondiagonal_item in nondiagonal_items:
            nondiagonal_item.append(chname)
            nondiagonal_ham_items.append(nondiagonal_item)
    
    if rh0 != 0:
        Ht.append(rh0*pi*2)   
    hamNO=0    
    for item in diagonal_ham_items:
        product,chname = item
        hamfunc='rotation_ham_' +str(hamNO)
        hamNO=hamNO+1
        output = StringIO.StringIO()
        output.write('def '+ hamfunc + '(t,args):\n')
        output.write('\tpulses,channels = args[\'PULS\'],args[\'MWCH\']\n')
        output.write('\tamp,phase = pulses.get_shape(\''+chname+'\',t)\n')
        
        output.write('\treturn cos(('+str(channels[chname].frequency)+\
                     '*t*pi*2+channels[\''+chname+'\'].phase + phase))  * amp\n')
        output.write('Ht.append([product*pi*2*'+str(channels[chname].field*2)+', '+hamfunc+'])\n')
        exec output.getvalue()
        print output.getvalue()
        output.close()
        max_freq = max(max_freq,abs(channels[chname].frequency))
        
    for item in nondiagonal_ham_items:
        product,omega_freq,chname = item
        max_freq = max(max_freq,abs(omega_freq))
        #print a,b,freq,sign
        hamfunc='rotation_ham_' +str(hamNO)
        hamNO=hamNO+1
        output = StringIO.StringIO()
        output.write('def '+ hamfunc + '(t,args):\n')
        if chname == 'static_item':
#             output.write('\tprint t\n')
            output.write('\treturn exp(1j*('+str(omega_freq)+'*t*pi*2))\n')
            output.write('Ht.append([product*pi, '+hamfunc+'])\n')
        else:      
            output.write('\tpulses,channels = args[\'PULS\'],args[\'MWCH\']\n')
            output.write('\tamp,phase = pulses.get_shape(\''+chname+'\',t)\n')   
#             output.write('\tprint amp\n')      
#             output.write('\tprint t,amp\n')
            output.write('\treturn exp(1j*('+str(omega_freq)+'*t*pi*2' + \
                     '+channels[\''+chname+'\'].phase + phase))  * amp\n')
            output.write('Ht.append([product*pi*'+str(channels[chname].field*2)+', '+hamfunc+'])\n')

        exec output.getvalue()
        if debug :
            print output.getvalue()
        hamfunc='rotation_ham_' +str(hamNO)
        hamNO=hamNO+1
        output = StringIO.StringIO()
        output.write('def '+ hamfunc + '(t,args):\n')

        if chname=='static_item' :
            output.write('\treturn exp(-1j*('+str(omega_freq)+'*t*pi*2))\n')
        else:      
            output.write('\tpulses,channels = args[\'PULS\'],args[\'MWCH\']\n')
            output.write('\tamp,phase = pulses.get_shape(\''+chname+'\',t)\n')         
            output.write('\treturn exp(-1j*('+str(omega_freq)+'*t*pi*2' + \
                     '+channels[\''+chname+'\'].phase + phase))  * amp\n')

        output.write('Ht.append([product.dag()*pi'+('*'+str(channels[chname].field*2) if chname!='static_item' else '')+', '+hamfunc+'])\n')
        if debug :
            print output.getvalue()
            
        exec output.getvalue()
        output.close()

    #options for mesolve
    opts=Odeoptions()
    opts.nsteps=1e6
    args={'MWCH':channels}    
    return max_freq,args,opts,Ht

def cal_propagator(args,Ht,trange,steps):        
    if not isinstance(trange,list):
        trange = [0,trange]
#     dt=(trange[1]-trange[0])/float(steps)#thresh_err/freq if abs(freq)> 0 else (trange[1]-trange[0])
#     nslices = int(floor((trange[1]-trange[0])/dt))
    dt = (trange[1]-trange[0])/float(steps)
    print 'dt = ',dt 
    #dt = float(trange[1] - trange[0]) / nslices 
    U = 1
    for n in range(int(steps)):
        t = trange[0] + n * dt
        
        ham = 0
        for n in range(0, len(Ht)):
            if isinstance(Ht[n],Qobj):
                ham = ham + Ht[n]
            else:
                ham = ham + Ht[n][0] * Ht[n][1](t, args) 
        U = (-1j*dt*ham).expm()*U
    return U      
    

def Bdx(t,args):
    channels = args['MWCH']
    pulses = args['PULS']
    b=0
    for ch in channels:
        ampl,phase = pulses.get_shape(ch,t)
        b+=2*channels[ch].field*ampl*cos(channels[ch].theta)*cos(channels[ch].phi)*cos(2*pi*channels[ch].frequency*t+phase + channels[ch].phase)  
    return b

def Bdy(t,args):
    channels = args['MWCH']
    pulses = args['PULS']
    b=0
    for ch in channels:
        ampl,phase = pulses.get_shape(ch,t)
        b+= 2*channels[ch].field*ampl*cos(channels[ch].theta)*sin(channels[ch].phi)*cos(2*pi*channels[ch].frequency*t+phase + channels[ch].phase)
    return b

    
def Bdz(t,args):
    channels = args['MWCH']
    pulses = args['PULS']
    b=0
    for ch in channels:
        ampl,phase = pulses.get_shape(ch,t)
        b+= 2*channels[ch].field*ampl*sin(channels[ch].theta)*cos(2*pi*channels[ch].frequency*t+2*pi + channels[ch].phase)
    return b

  


    